/*
 * =====================================================================================
 *
 *       Filename:  gmm.h
 *
 *    Description:  gaussian mixture model
 *
 *        Version:  1.0
 *        Created:  2009年06月15日 21时08分06秒
 *       Revision:  none
 *       Compiler:  gcc
 *
 *         Author:  Ying Wang (WY), ywang@nlpr.ia.ac.cn
 *        Company:  Institute of Automation, Chinese Academy of Sciences
 *
 * =====================================================================================
 */
#ifndef GMM_H
#define GMM_H
#include "statistical.h"
#include "kmeans.h"
#define GMM_ITER_NUM 60
#define GMM_THRESHOLD 1e-8
class preGMM
{
public:
    class Mat_mm : public NCmatrix<double>
    {
    public:
        Mat_mm()
            :NCmatrix<double> (mmstat,mmstat)
        {
        }

    };
    preGMM(int m)
    {
        mmstat= m;
    }
    static int mmstat;

};
int preGMM::mmstat = -1;


/**
 * \brief Gaussian mixture model
 *  useage: GMM mix(matrix,K,Covtype); means=mix.means; covs=mix.covs; resp = mix.resp;
 */
class GMM : public preGMM
{
public:
     GMM(const NCmatrix<double> &mdata, int K, Covtype covtype=Full)
        : preGMM(mdata.column())
        , data(mdata)
        , KK(K)
        , N(mdata.row())
        , dim(mdata.column())
        , type(covtype)
        , means(KK,dim)
        , fracs(KK)
        , resp(N,KK)
        , lndets(KK)
        , covs(KK)
    {
        for(int k=0;k<KK; k++ )
        {
            fracs[k] = 1./KK;
            for(int i=0;i<dim;i++)
            {
                for(int j=0;j<dim;j++)
                {
                    covs[k][i][j] =0.;
                }
                covs[k][i][i] = 1.0;
            }
        }
    }
public:

	void fit();
private:
	double estep();
	void mstep();
public:
    	NCmatrix<double> data;
	int KK;
	int N;
	int dim;
	Covtype type;
	NCmatrix<double> means;
	NCvector<double> fracs;
	NCmatrix<double> resp;
	NCvector<double> lndets;
	std::vector<Mat_mm > covs;
};

void GMM::fit()
{
	double oldloglike, loglike,change_percent;
	int iter=0;
	//	std::cout<<"here"<<std::endl;
	kmeans(data,KK,means);
	oldloglike = estep();
	mstep();

	do
	{
		loglike = estep();
		mstep();
		iter++;
		change_percent = -(loglike-oldloglike)/oldloglike;
		oldloglike= loglike;
		std::cout<<"loglike: "<< oldloglike<<";"<<"iter num: " <<iter<<";"<<"change_percent: " << change_percent <<std::endl;

	}while(iter<GMM_ITER_NUM && fabs(change_percent)>GMM_THRESHOLD);

	std::cout<<"----------"<<std::endl;

}

double GMM::estep()
{
	int k,n;
	double sum=0,tmploglike=0,maxnum,tmp;

	NCvector<double> x,imean;


	for( n=0; n<N; n++ )
	{
	//	x = data.rowvector(n);
		for( k=0; k<KK; k++ )
		{
	//		imean = means.rowvector(k);
			resp[n][k] = log(fracs[k])+loggausspdf(data.rowvector(n),means.rowvector(k),covs[k],type);
		}
	}

   	//std::cout<<"e s" <<std::endl;
	for( n=0; n<N; n++ )
	{
		for( k=0; k<KK; k++ ) if (resp[n][k] > maxnum ) maxnum = resp[n][k];
		for(sum=0.,k=0;k<KK; k++ ) sum += exp(resp[n][k] - maxnum);
		tmp = maxnum+log(sum);
		for( k=0; k<KK; k++ ) {resp[n][k] = exp(resp[n][k] - tmp);}
		tmploglike += tmp;
	}
	return tmploglike;
}

void GMM::mstep()
{
	int n,k,j,jj;
	double Ni,sum1,sum2;
	if(type == Full)
	{
 		for( k=0; k<KK; k++ )
 		{
 			Ni=0.;
 			for( n=0; n<N; n++ )
 				Ni += resp[n][k];
 			fracs[k] = Ni/N;
 			for( j=0; j<dim; j++ )
 			{
 				for(sum1=0.,n=0; n<N; n++ )
 					sum1 += resp[n][k] *data[n][j];
 				means[k][j] = sum1/Ni;
				for ( jj=0; jj<dim; jj++ )
				{
					for( sum2=0., n=0; n<N; n++ )
					{
						sum2 += resp[n][k]* (data[n][j]-means[k][j])*(data[n][jj]-means[k][jj]);
					}
					covs[k][j][jj] = sum2/Ni;

				}
 			}
 		}
	}
	else
	{
		for( k=0; k<KK; k++ )
		{
			Ni=0.;
			for( n=0; n<N; n++ )
				Ni += resp[n][k];
			fracs[k] = Ni/N;
			for( j=0; j<dim; j++ )
			{
				for( sum1=0.,sum2=0., n=0; n<N; n++ )
				{
					sum1 += resp[n][k] * data[n][j];
					sum2 += resp[n][k] * SQE(data[n][j]-means[k][j]);
				}
				means[k][j] = sum1/Ni;
				covs[k][j][j] = sum2/Ni;
			}
		}
	}
}
#endif
